"""
更新MongoDB中热点数据的时间戳脚本
将所有热点数据的时间更新为当前时间，使其在前端显示为最新数据
"""
import os
import sys
from datetime import datetime, timedelta
from pymongo import MongoClient
from typing import List

# 添加项目路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def get_mongo_client():
    """获取MongoDB客户端"""
    # 从环境变量获取，或使用默认值（带认证）
    mongo_uri = os.getenv('MONGODB_URI', 'mongodb://admin:newhand@localhost:27017/')
    client = MongoClient(mongo_uri)
    return client


def update_timestamps_batch(collection, platform: str = None, days_ago: int = 0):
    """
    批量更新时间戳

    Args:
        collection: MongoDB collection
        platform: 指定平台，None表示所有平台
        days_ago: 几天前的数据（0表示当前时间）
    """
    # 构建查询条件
    query = {}
    if platform:
        query['platform'] = platform

    # 计算新的时间戳
    new_time = datetime.now() - timedelta(days=days_ago)

    # 更新操作
    update_fields = {
        '$set': {
            'crawled_at': new_time,
            'updated_at': new_time,
            'publish_time': new_time
        }
    }

    # 执行批量更新
    result = collection.update_many(query, update_fields)

    return result.modified_count


def update_with_rank_distribution(collection, platform: str = None):
    """
    按排名分布更新时间，让排名靠前的显示更新

    Args:
        collection: MongoDB collection
        platform: 指定平台
    """
    query = {}
    if platform:
        query['platform'] = platform

    # 获取所有数据
    items = list(collection.find(query).sort('rank', 1))

    updated_count = 0
    now = datetime.now()

    for idx, item in enumerate(items):
        # 排名越靠前，时间越新
        # 前10名：当前时间
        # 11-30名：30分钟前
        # 31-50名：1小时前
        # 其他：2小时前
        rank = item.get('rank', idx + 1)

        if rank <= 10:
            time_delta = timedelta(minutes=0)
        elif rank <= 30:
            time_delta = timedelta(minutes=30)
        elif rank <= 50:
            time_delta = timedelta(hours=1)
        else:
            time_delta = timedelta(hours=2)

        new_time = now - time_delta

        # 更新单条记录
        collection.update_one(
            {'_id': item['_id']},
            {
                '$set': {
                    'crawled_at': new_time,
                    'updated_at': new_time,
                    'publish_time': new_time
                }
            }
        )
        updated_count += 1

    return updated_count


def main():
    """主函数"""
    print("=" * 60)
    print("MongoDB热点数据时间戳更新工具")
    print("=" * 60)

    # 连接MongoDB
    try:
        client = get_mongo_client()
        db_name = os.getenv('MONGODB_DB', 'crawler_db')  # 从环境变量获取数据库名
        db = client[db_name]
        collection = db['hot_items']  # 集合名称

        print(f"\n已连接到MongoDB")
        print(f"数据库: {db.name}")
        print(f"集合: {collection.name}")

        # 显示当前数据统计
        total_count = collection.count_documents({})
        print(f"\n当前总数据量: {total_count}")

        # 按平台统计
        platforms = collection.distinct('platform')
        print(f"平台列表: {', '.join(platforms)}")

        for platform in platforms:
            count = collection.count_documents({'platform': platform})
            print(f"  - {platform}: {count} 条")

        # 选择更新模式
        print("\n" + "=" * 60)
        print("请选择更新模式:")
        print("1. 全部更新为当前时间")
        print("2. 按排名分布更新时间（推荐）")
        print("3. 更新指定平台")
        print("4. 退出")
        print("=" * 60)

        choice = input("\n请输入选项 (1-4): ").strip()

        if choice == '1':
            print("\n正在更新所有数据...")
            count = update_timestamps_batch(collection)
            print(f"✅ 成功更新 {count} 条数据")

        elif choice == '2':
            print("\n正在按排名分布更新时间...")
            total_updated = 0
            for platform in platforms:
                print(f"\n更新平台: {platform}")
                count = update_with_rank_distribution(collection, platform)
                print(f"  ✅ 更新 {count} 条")
                total_updated += count
            print(f"\n总计更新: {total_updated} 条数据")

        elif choice == '3':
            print(f"\n可用平台: {', '.join(platforms)}")
            platform = input("请输入平台名称: ").strip().lower()

            if platform in platforms:
                print(f"\n正在更新平台 {platform}...")
                count = update_timestamps_batch(collection, platform)
                print(f"✅ 成功更新 {count} 条数据")
            else:
                print(f"❌ 无效的平台名称: {platform}")

        elif choice == '4':
            print("\n已退出")
            return
        else:
            print("\n❌ 无效的选项")
            return

        # 显示更新后的样例数据
        print("\n" + "=" * 60)
        print("更新后的样例数据:")
        print("=" * 60)

        sample = collection.find_one()
        if sample:
            print(f"标题: {sample.get('title', 'N/A')}")
            print(f"平台: {sample.get('platform', 'N/A')}")
            print(f"排名: {sample.get('rank', 'N/A')}")
            print(f"爬取时间: {sample.get('crawled_at', 'N/A')}")
            print(f"更新时间: {sample.get('updated_at', 'N/A')}")
            print(f"发布时间: {sample.get('publish_time', 'N/A')}")

        print("\n✅ 时间戳更新完成！")
        print("现在前端应该可以看到最新的热点数据了")

    except Exception as e:
        print(f"\n❌ 错误: {str(e)}")
        import traceback
        traceback.print_exc()
    finally:
        if 'client' in locals():
            client.close()
            print("\nMongoDB连接已关闭")


if __name__ == '__main__':
    main()
